from math import floor
from re import S
import cv2
import os
import numpy as np
import glob
import torch
import tqdm
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms

mean = [0.379, 0.398, 0.384]
std = [0.311, 0.319, 0.329]

# 训练和测试的数据集形式上有区别,测试只需要传入图片,同时展示的时候需要在原图上画,所以转tensor需要手动

my_kitti_path = '/Dataset/KITTI_SEG/training/'

class KITTISegDataset(Dataset):
    def __init__(self, kitti_path) -> None:
        super().__init__()
        self.images = glob.glob(kitti_path + "image_2/*.png")
        self.masks = glob.glob(kitti_path + "semantic/*.png")
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=mean, std=std)
        ])
        
    def __getitem__(self, i):
        assert self.images[i].split(".png")[0].split("/")[-1] == self.masks[i].split(".png")[0].split("/")[-1]
        img = cv2.imread(self.images[i])
        mask = cv2.imread(self.masks[i], 0)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        img = self.transform(img)
        mask = torch.tensor(mask)
        rect = transforms.RandomCrop.get_params(img, ((352, 1216)))
        img = transforms.functional.crop(img, *rect)
        mask = transforms.functional.crop(mask, *rect)
        
        return img, mask
    
    def __len__(self):
        return len(self.images)

# 测试用
class KITTI_IMG(Dataset):
    def __init__(self) -> None:
        super().__init__()
        # self.images = glob.glob(kitti_path + "image_2/*.png")
        self.images = glob.glob("/Dataset/KITTI_ODMO/sequences/00/image_2/*.png")
        self.images.sort()
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=mean, std=std)
        ])
        self.imgsize = cv2.imread(self.images[0]).shape
        # 裁剪后的尺寸
        self.h = floor(self.imgsize[0] / 32) * 32
        self.w = floor(self.imgsize[1] / 32) * 32
    
    def __getitem__(self, i):
        img = cv2.imread(self.images[i])
        # 先裁剪,保证输出图与原图保持一致
        img = img[:self.h, :self.w, :]
        # print(img.shape)
        cvt_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        return img, self.transform(cvt_img)
    
    def __len__(self):
        return len(self.images)


def loadkitti(kitti_path):
    return KITTISegDataset(kitti_path)

# 计算图像数据的均值和标准差
def compute_mean_std(kitti_path):
    data = loadkitti(kitti_path)
    mean_r = 0
    mean_g = 0
    mean_b = 0
    
    for i in tqdm.tqdm(data):
        mean_r += np.mean(i[0][:, :, 0])
        mean_g += np.mean(i[0][:, :, 1])
        mean_b += np.mean(i[0][:, :, 2])
        
    mean_r /= len(data)
    mean_g /= len(data)
    mean_b /= len(data)
    
    diff_r = 0
    diff_g = 0
    diff_b = 0
    N = 0
    
    for i in tqdm.tqdm(data):
        diff_r += np.sum(np.power(i[0][:, :, 0] - mean_r, 2))
        diff_g += np.sum(np.power(i[0][:, :, 1] - mean_g, 2))
        diff_b += np.sum(np.power(i[0][:, :, 2] - mean_b, 2))
        N += np.prod(i[0][:, :, 0].shape)
        
    std_r = np.sqrt(diff_r / N)
    std_g = np.sqrt(diff_g / N)
    std_b = np.sqrt(diff_b / N)
    
    mean = [mean_r.item() / 255.0, mean_g.item() / 255.0, mean_b.item() / 255.0]
    std = [std_r.item() / 255.0, std_g.item() / 255.0, std_b.item() / 255.0]

    return mean, std
    

if __name__ == "__main__":
    # KITTI = loadkitti(my_kitti_path)
    KITTI = KITTI_IMG()
    # data = DataLoader(KITTI)
    for i in KITTI:
        print(i[1])
        break
    